from openai import OpenAI  # Tool scripts do not need to be run directly.
# noinspection PyUnresolvedReferences
import openai
# noinspection PyUnresolvedReferences
import numpy as np
import time

# Set your api key
my_api_key = "YOUR_API_KEY"
client = OpenAI(
    # defaults to os.environ.get("OPENAI_API_KEY")
    api_key = my_api_key,
    timeout=20.0,

)



# def do_query(system_prompt, max_tokens=2, engine="gpt-3.5-turbo"):
#     messages = [{"role": "system", "content": system_prompt}]
#     return system_prompt


def do_query(system_prompt, user_prompt, max_tokens=2, engine="gpt-3.5-turbo-0613"):
    messages = [{"role": "system", "content": system_prompt}, {"role": "user", "content": user_prompt}]
    response = client.chat.completions.create(
        model=engine,
        messages=messages,
        max_tokens=max_tokens,
    )
    return response.choices[0].message.content

def run_prompts(prompts, engine="gpt-3.5-turbo-0613"):
    results = []
    for prompt in prompts:
        system_prompt = do_query(prompt, max_tokens=2, engine=engine)
        results.append(system_prompt)
        time.sleep(0.1)
    return results

